/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.transform;



import com.bff.gaia.unified.sdk.extensions.sql.UnifiedSqlSeekableTable;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.SerializableRexFieldAccess;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.SerializableRexInputRef;

import com.bff.gaia.unified.sdk.extensions.sql.impl.utils.SerializableRexNode;

import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.transforms.SimpleFunction;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.Row;

import org.apache.calcite.rel.core.JoinRelType;

import org.apache.calcite.rex.RexCall;

import org.apache.calcite.rex.RexInputRef;

import org.apache.calcite.rex.RexNode;

import org.apache.calcite.util.Pair;



import java.util.ArrayList;

import java.util.List;

import java.util.Map;



import static java.util.stream.Collectors.toList;



/** Collections of {@code PTransform} and {@code DoFn} used to perform JOIN operation. */

public class UnifiedJoinTransforms {



  /** A {@code SimpleFunction} to extract join fields from the specified row. */

  public static class ExtractJoinFields extends SimpleFunction<Row, KV<Row, Row>> {

    private final List<SerializableRexNode> joinColumns;

    private final Schema schema;

    private int leftRowColumnCount;



    public ExtractJoinFields(

        boolean isLeft,

        List<Pair<RexNode, RexNode>> joinColumns,

        Schema schema,

        int leftRowColumnCount) {

      this.joinColumns =

          joinColumns.stream()

              .map(pair -> SerializableRexNode.builder(isLeft ? pair.left : pair.right).build())

              .collect(toList());

      this.schema = schema;

      this.leftRowColumnCount = leftRowColumnCount;

    }



    @Override

    public KV<Row, Row> apply(Row input) {

      Row row =

          joinColumns.stream()

              .map(v -> getValue(v, input, leftRowColumnCount))

              .collect(Row.toRow(schema));

      return KV.of(row, input);

    }



    @SuppressWarnings("unused")

    private Schema.Field toField(Schema schema, Integer fieldIndex) {

      Schema.Field original = schema.getField(fieldIndex);

      return original.withName("c" + fieldIndex);

    }



    private Object getValue(

		SerializableRexNode serializableRexNode, Row input, int leftRowColumnCount) {

      if (serializableRexNode instanceof SerializableRexInputRef) {

        return input.getValue(

            ((SerializableRexInputRef) serializableRexNode).getIndex() - leftRowColumnCount);

      } else { // It can only be SerializableFieldAccess.

        List<Integer> indexes = ((SerializableRexFieldAccess) serializableRexNode).getIndexes();

        // retrieve row based on the first column reference.

        Row rowField = input.getValue(indexes.get(0) - leftRowColumnCount);

        for (int i = 1; i < indexes.size() - 1; i++) {

          rowField = rowField.getRow(indexes.get(i));

        }

        return rowField.getValue(indexes.get(indexes.size() - 1));

      }

    }

  }



  /** A {@code DoFn} which implement the sideInput-JOIN. */

  public static class SideInputJoinDoFn extends DoFn<KV<Row, Row>, Row> {

    private final PCollectionView<Map<Row, Iterable<Row>>> sideInputView;

    private final JoinRelType joinType;

    private final Row rightNullRow;

    private final boolean swap;

    private final Schema schema;



    public SideInputJoinDoFn(

        JoinRelType joinType,

        Row rightNullRow,

        PCollectionView<Map<Row, Iterable<Row>>> sideInputView,

        boolean swap,

        Schema schema) {

      this.joinType = joinType;

      this.rightNullRow = rightNullRow;

      this.sideInputView = sideInputView;

      this.swap = swap;

      this.schema = schema;

    }



    @ProcessElement

    public void processElement(ProcessContext context) {

      Row key = context.element().getKey();

      Row leftRow = context.element().getValue();

      Map<Row, Iterable<Row>> key2Rows = context.sideInput(sideInputView);

      Iterable<Row> rightRowsIterable = key2Rows.get(key);



      if (rightRowsIterable != null && rightRowsIterable.iterator().hasNext()) {

        for (Row aRightRowsIterable : rightRowsIterable) {

          context.output(combineTwoRowsIntoOne(leftRow, aRightRowsIterable, swap, schema));

        }

      } else {

        if (joinType == JoinRelType.LEFT) {

          context.output(combineTwoRowsIntoOne(leftRow, rightNullRow, swap, schema));

        }

      }

    }

  }



  /** A {@code SimpleFunction} to combine two rows into one. */

  public static class JoinParts2WholeRow extends SimpleFunction<KV<Row, KV<Row, Row>>, Row> {

    private final Schema schema;



    public JoinParts2WholeRow(Schema schema) {

      this.schema = schema;

    }



    @Override

    public Row apply(KV<Row, KV<Row, Row>> input) {

      KV<Row, Row> parts = input.getValue();

      Row leftRow = parts.getKey();

      Row rightRow = parts.getValue();

      return combineTwoRowsIntoOne(leftRow, rightRow, false, schema);

    }

  }



  /** As the method name suggests: combine two rows into one wide row. */

  private static Row combineTwoRowsIntoOne(

	  Row leftRow, Row rightRow, boolean swap, Schema outputSchema) {

    if (swap) {

      return combineTwoRowsIntoOneHelper(rightRow, leftRow, outputSchema);

    } else {

      return combineTwoRowsIntoOneHelper(leftRow, rightRow, outputSchema);

    }

  }



  /** As the method name suggests: combine two rows into one wide row. */

  private static Row combineTwoRowsIntoOneHelper(Row leftRow, Row rightRow, Schema ouputSchema) {

    return Row.withSchema(ouputSchema)

        .addValues(leftRow.getValues())

        .addValues(rightRow.getValues())

        .build();

  }



  /** Transform to execute Join as Lookup. */

  public static class JoinAsLookup extends PTransform<PCollection<Row>, PCollection<Row>> {

    private final UnifiedSqlSeekableTable seekableTable;

    private final Schema lkpSchema;

    private final int factColOffset;

    private Schema joinSubsetType;

    private final Schema outputSchema;

    private List<Integer> factJoinIdx;



    public JoinAsLookup(

        RexNode joinCondition,

        UnifiedSqlSeekableTable seekableTable,

        Schema lkpSchema,

        Schema outputSchema,

        int factColOffset,

        int lkpColOffset) {

      this.seekableTable = seekableTable;

      this.lkpSchema = lkpSchema;

      this.outputSchema = outputSchema;

      this.factColOffset = factColOffset;

      joinFieldsMapping(joinCondition, factColOffset, lkpColOffset);

    }



    private void joinFieldsMapping(RexNode joinCondition, int factColOffset, int lkpColOffset) {

      factJoinIdx = new ArrayList<>();

      List<Schema.Field> lkpJoinFields = new ArrayList<>();



      RexCall call = (RexCall) joinCondition;

      if ("AND".equals(call.getOperator().getName())) {

        List<RexNode> operands = call.getOperands();

        for (RexNode rexNode : operands) {

          factJoinIdx.add(

              ((RexInputRef) ((RexCall) rexNode).getOperands().get(0)).getIndex() - factColOffset);

          int lkpJoinIdx =

              ((RexInputRef) ((RexCall) rexNode).getOperands().get(1)).getIndex() - lkpColOffset;

          lkpJoinFields.add(lkpSchema.getField(lkpJoinIdx));

        }

      } else if ("=".equals(call.getOperator().getName())) {

        factJoinIdx.add(((RexInputRef) call.getOperands().get(0)).getIndex() - factColOffset);

        int lkpJoinIdx = ((RexInputRef) call.getOperands().get(1)).getIndex() - lkpColOffset;

        lkpJoinFields.add(lkpSchema.getField(lkpJoinIdx));

      } else {

        throw new UnsupportedOperationException(

            "Operator " + call.getOperator().getName() + " is not supported in join condition");

      }



      joinSubsetType = Schema.builder().addFields(lkpJoinFields).build();

    }



    @Override

    public PCollection<Row> expand(PCollection<Row> input) {

      return input

          .apply(

              "join_as_lookup",

              ParDo.of(

                  new DoFn<Row, Row>() {

                    @Setup

                    public void setup() {

                      seekableTable.setUp();

                    }



                    @ProcessElement

                    public void processElement(ProcessContext context) {

                      Row factRow = context.element();

                      Row joinSubRow = extractJoinSubRow(factRow);

                      List<Row> lookupRows = seekableTable.seekRow(joinSubRow);

                      for (Row lr : lookupRows) {

                        context.output(

                            combineTwoRowsIntoOne(factRow, lr, factColOffset != 0, outputSchema));

                      }

                    }



                    @Teardown

                    public void teardown() {

                      seekableTable.tearDown();

                    }



                    private Row extractJoinSubRow(Row factRow) {

                      List<Object> joinSubsetValues =

                          factJoinIdx.stream().map(factRow::getValue).collect(toList());



                      return Row.withSchema(joinSubsetType).addValues(joinSubsetValues).build();

                    }

                  }))

          .setRowSchema(joinSubsetType);

    }

  }

}